import getpass
import os
from dotenv import load_dotenv; load_dotenv()

from typing import Annotated

from typing_extensions import TypedDict

from langgraph.graph import StateGraph, START, END
from langgraph.graph.message import add_messages

from langchain_openai import ChatOpenAI

from IPython.display import Image, display


def _set_env(var: str):
    if not os.environ.get(var):
        os.environ[var] = getpass.getpass(f"{var}: ")

class State(TypedDict):
    # Messages have the type "list". The `add_messages` function
    # in the annotation defines how this state key should be updated
    # (in this case, it appends messages to the list, rather than overwriting them)
    messages: Annotated[list, add_messages]


def chatbot(state: State):
    message = llm.invoke(state["messages"])
    return {"messages": [message]}


if __name__ == "__main__":
	_set_env("BAICHUAN_API_KEY")

	graph_builder = StateGraph(State)

	# llm = ChatAnthropic(model="claude-3-haiku-20240307")
	llm = ChatOpenAI(model="Baichuan2", api_key="123", base_url="http://bc.192.168.107.2.nip.io/v1")

	# The first argument is the unique node name
	# The second argument is the function or object that will be called whenever
	# the node is used.
	graph_builder.add_node("chatbot", chatbot)
	graph_builder.add_edge(START, "chatbot")
	graph_builder.add_edge("chatbot", END)
	graph = graph_builder.compile()
 
	try:
		# 提取文件名前缀，不包含目录
		graph.get_graph().draw_mermaid_png(output_file_path="./" + os.path.basename(__file__).split(".")[0] + "_graph.png")
		# display(Image(graph.get_graph().draw_mermaid_png()))
	except Exception:
		# This requires some extra dependencies and is optional
		print(f"Error: {Exception}")
		pass

	while True:
		user_input = input("User: ")
		if user_input.lower() in ["quit", "exit", "q"]:
			print("Goodbye!")
			break
		for event in graph.stream({"messages": ("user", user_input)}):
			for value in event.values():
				print("Assistant:", value["messages"][-1].content)

